#! /usr/bin/env python

import numpy as np
import tensorflow as tf
from tensorflow.python.framework import graph_util


input_nodes = ['input', 'keep_prob']    # your model input
output_nodes = ['output']   # what you want to get, must be string list


def freeze_graph(path='model.ckpt', output='model.pb'):
    saver = tf.train.import_meta_graph(path+'.meta', clear_devices=True)
    graph = tf.get_default_graph()
    input_graph_def = graph.as_graph_def()

    with tf.Session() as sess:
        saver.restore(sess, path)
        output_graph_def = graph_util.convert_variables_to_constants(
                           sess=sess,
                           input_graph_def=input_graph_def,   # = sess.graph_def,
                           output_node_names=output_nodes)

        with tf.gfile.GFile(output, 'wb') as fgraph:
            fgraph.write(output_graph_def.SerializeToString())


def freeze_single(path='model.pb'):
    with tf.gfile.GFile(path, 'rb') as fgraph:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(fgraph.read())

    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
        sess = tf.Session(graph=graph)

        inputs = [graph.get_tensor_by_name(n+':0') for n in input_nodes]
        output = [graph.get_tensor_by_name(n+':0') for n in output_nodes]


        return sess, inputs, output


if __name__ == '__main__':
    freeze_graph(path='model.ckpt', output='model.pb')

    sess, inputs, output = freeze_single('model.pb')
    feed_dict = {n: d for n, d in zip(inputs, [np.random.randn(), 1.0])}
    result = sess.run(output, feed_dict=feed_dict)
    print(result)
